import struct

from PyQt5.Qt import *

from ExtendWidgets.notifycation import NotificationWindow
from UI.dockwidget import Ui_DockWidget
import pyqtgraph as pg
from pyqtgraph.graphicsItems.PlotItem import PlotItem
import numpy as np
import os
import shutil
from model_train.train_model import TrainProcess, Args
import cgitb
import json

cgitb.enable(format='text')
pg.setConfigOptions(antialias=True)
import scipy
from scipy.signal import stft
from typing import *

html_str = """
<style type="text/css">
table.gridtable {
    font-family: verdana,arial,sans-serif;
    font-size:13px;
    color:#333333;
    border-width: 1px;
    border-color: #666666;
    border-collapse: collapse;

}
table.gridtable th {
    border-width: 1px;
    padding: 8px;
    border-style: solid;
    border-color: #666666;
    background-color: #dedede;
}
table.gridtable td {
    border-width: 1px;
    padding: 8px;
    border-style: solid;
    border-color: #666666;
    background-color: #ffffff;
    text-align: center;
}
</style>
<body>
<table class="gridtable">
"""


def GetHtml(header, data):
    content = """"""
    columns = header
    content += "<tr>"
    for title in columns:
        content += "<th>{}</th>".format(title)
    content += "</tr>"
    for i in range(len(data)):
        content += "<tr>"
        for j in range(len(data[0])):
            content += '<td><b>{}</b></td>'.format(data[i][j])
        content += "</tr>"
    content += "</table></body>"
    return html_str + content


index = 0


# 读取文件线程
class BinFileReader(QThread):
    requestShowProcess = pyqtSignal(int)
    requestShowFinished = pyqtSignal(str)
    requestPassTrainData = pyqtSignal(object)

    def __init__(self, parent, train_cnt, bin_file_path, noverlap, time_segment=10, nfft=2048, Fs=100e6, nperseg=None,
                 F0=None, begin=0, end=None):
        super(BinFileReader, self).__init__(parent)
        global index, flag
        index = 0
        flag = False
        self.mutex = QMutex()
        self.nperseg = nfft if nperseg is None else nperseg
        self.train_cnt = train_cnt
        self.parent = parent
        self.bin_file_path = bin_file_path
        self.fp = None
        self.nfft = nfft
        self.Fs = Fs
        self.F0 = F0
        self.noverlap = noverlap
        self.width = nfft
        self.time_segment = time_segment
        self.iq_complex_cnt = self.getReadLength(self.time_segment, self.nfft, self.noverlap)
        self.begin = begin
        self.end = end

    def openFile(self):
        try:
            self.fp = open(self.bin_file_path, "rb")  # 以二进制只读方式打开文件
            return True
        except OSError:
            return False

    def getReadLength(self, time_col, window_length, noverlap):
        return (time_col - 3) * (window_length - noverlap) + window_length

    def run(self):
        print("进入文件读取线程")
        try:
            ret = self.openFile()
            if not ret:
                QMessageBox.critical(self, "错误提示", "打开文件失败")
                return
        except:
            QMessageBox.critical(self, "错误提示", "打开文件失败")
            print("打开测试文件失败")
            return
        # 读取文件
        if self.isInterruptionRequested():
            self.parent.running = False
            return
        # train_data = np.zeros((self.train_cnt, self.time_segment, self.nfft), dtype=np.float)  # 预先分配内存
        train_data = np.zeros((self.train_cnt, self.time_segment, self.end - self.begin + 1), dtype=np.float)  # 预先分配内存

        iq_complex_cnt = self.iq_complex_cnt
        buffer_num = 1
        read_num = iq_complex_cnt * 2  # 生成一张时頻图需要的IQ点数(非复数)
        read_bin_cnt = read_num * 2 * buffer_num

        I_index = np.arange(0, read_num, 2)
        Q_index = np.arange(1, read_num, 2)
        nfft = self.nfft
        nperseg = self.nperseg
        window = scipy.signal.get_window('hamming', Nx=nperseg)
        Fs = self.Fs
        noverlap = self.noverlap
        print("数据加载完毕")

        def IQ_to_PSD(IQ):
            global flag, index
            length = IQ.shape[0]
            for i in range(length):
                if self.isInterruptionRequested():
                    return
                I = IQ[i][I_index]
                Q = IQ[i][Q_index]
                IQ_complex = np.zeros([iq_complex_cnt], dtype=np.complex)
                IQ_complex.real = I
                IQ_complex.imag = Q
                F, T, Zxx = stft(IQ_complex, fs=Fs, noverlap=noverlap, nfft=nfft, window=window,
                                 return_onesided=False,
                                 nperseg=nperseg)
                s: np.ndarray = np.abs(Zxx) * 2 / nfft
                # s = (s - s.min()) / (s.max() - s.min())
                s = s / s.std()

                train_data[index] = s.T[:, self.begin:self.end + 1]
                index += 1

        while index < self.train_cnt:
            try:
                self.mutex.lock()
                if self.isInterruptionRequested():
                    self.parent.running = False
                    print("退出程序")
                    return
                content = self.fp.read(read_bin_cnt)
                IQ = struct.unpack('h' * read_num * buffer_num, content)
                IQ = np.array(IQ).reshape([-1, read_num])
                IQ_to_PSD(IQ)  # IQ数据转PSD并保存到文件队列
                self.mutex.unlock()
                self.requestShowProcess.emit(index)
            except Exception as e:
                print(e)
        self.requestPassTrainData.emit(train_data)


class DockWidget(Ui_DockWidget, QDockWidget):
    windowClosed = pyqtSignal()
    requestPausedThread = pyqtSignal()
    requestSaveCurrentModel = pyqtSignal()
    requestRecoverThread = pyqtSignal()

    def __init__(self, parent=None, n=-1, name="", band="", stftsetting: Dict[str, float] = None, **kwargs):

        super(DockWidget, self).__init__(parent)
        self.setupUi(self)
        self.stftsetting = stftsetting
        print(self.stftsetting)
        self.parent = parent

        from main import MainWin
        self.parent: MainWin = parent

        # todo: remove  可能死锁

        from windows.settingdockwidget import SettingDockWidget
        self.global_setting_dock: SettingDockWidget = self.parent.setting_dock
        self.band = band
        # self.global_setting_dock  = self.parent.setting_dock
        self.widget_rec.setPredefinedGradient('spectrum')
        self.widget_input.setPredefinedGradient('spectrum')
        self.MultiTrain = False
        self.printToTextBrowser("******************************训练SAIFE模型进行异常检测******************************")
        # 设置类选择器, 供设置样式
        self.pushButton_get_train_data_path.setStyleSheet("min-width: 25px")
        self.pushButton_get_picture_save_path.setStyleSheet("min-width: 25px")
        self.pushButton_get_model_save_path.setStyleSheet("min-width: 25px")

        # 3个设置
        self.global_path_setting = self.global_setting_dock.path_setting
        self.global_freq_range_setting = self.parent.setting_dock.train_setting
        self.global_threshold_setting = self.parent.setting_dock.model_manage

        self.training = False
        self.test_loss_lst = []
        self.train_loss_lst = []
        self.splitter.setStretchFactor(0, 2)
        self.splitter.setStretchFactor(1, 8)
        self.splitter.setStretchFactor(2, 4)
        self.lineEdit_model_name.setReadOnly(True)
        self.splitter_2.setStretchFactor(0, 3)
        self.splitter_2.setStretchFactor(1, 4)
        self.plotLossItem: PlotItem = self.plot_loss_widget.plotItem
        self.plotLossItem.setLabel(axis='bottom', text="训练轮数")
        self.plotLossItem.setLabel(axis='left', text="损失")
        self.plotLossItem.setTitle('重构损失变化曲线')
        self.plot_loss_widget.addLegend()
        self.dev_plot = self.plotLossItem.plot(pen='g', name='验证集损失')
        self.train_plot = self.plotLossItem.plot(pen='r', name="训练集损失")

        self.name = name
        self.n = n

        # 进度条设置
        self.progress_bar.setFormat("当前进度:%.2f%%" % 0)
        self.progress_bar.setAlignment(Qt.AlignVCenter | Qt.AlignRight)
        self.progress_bar.setMinimum(0)
        self.progress_bar.setValue(0)

        self.waiting = False
        self.lineEdit_train_data_path.setToolTip(self.lineEdit_train_data_path.text())
        self.lineEdit_picture_save_path.setToolTip(self.lineEdit_picture_save_path.text())

        self.lineEdit_model_save_path.cursorPositionChanged.connect(
            lambda: self.lineEdit_model_save_path.setCursorPosition(0))
        self.lineEdit_train_data_path.cursorPositionChanged.connect(
            lambda: self.lineEdit_train_data_path.setCursorPosition(0))
        self.lineEdit_picture_save_path.cursorPositionChanged.connect(
            lambda: self.lineEdit_picture_save_path.setCursorPosition(0))

        self.test()

    def test(self):
        self.spinBox_epoch.setValue(10)
        # self.spinBox_dimz.setValue(50)
        self.doubleSpinBox_LR.setValue(0.0005)
        self.spinBox_batch_size.setValue(32)
        self.spinBox_train_num.setValue(1000)
        self.lineEdit_model_save_path.setText(r'D:\test\model')
        self.lineEdit_train_data_path.setText(r"D:\异常检测项目最新版\DATA\DATA5yi.bin")
        self.lineEdit_picture_save_path.setText(r'D:\test\pictures')
        self.lineEdit_model_save_path.setCursorPosition(0)
        self.lineEdit_train_data_path.setCursorPosition(0)
        self.lineEdit_picture_save_path.setCursorPosition(0)

    def closeEvent(self, event: QCloseEvent) -> None:
        try:
            if self.training:
                ret = QMessageBox.warning(self, "窗口关闭警告", "关闭窗口后当前窗口执行的程序将被关闭，确定要关闭窗口吗？",
                                          QMessageBox.Ok | QMessageBox.Cancel)
                # ret = QMessageBox.warning(self, "窗口关闭警告", "当前正在训练模型, 不能关闭窗口",
                #                           QMessageBox.Ok )
                # event.setAccepted(False)
                if ret == QMessageBox.Cancel:
                    event.setAccepted(False)
                    return
                else:
                    self.thread.requestInterruption()
                    address = self.parent.findChild(DockWidget, name=self.objectName())
                    self.parent.docks.remove(address)
                    self.parent.needToDeleteDocks.append(address)
            else:
                self.parent.docks.remove(self.parent.findChild(DockWidget, name=self.objectName()))
                self.deleteLater()
        except Exception as e:
            print(e)

    def updateDealDataProcessbar(self, x):
        self.progress_bar.setValue(x)
        p = (self.progress_bar.value() - self.progress_bar.minimum()) * 100 / (
                self.progress_bar.maximum() - self.progress_bar.minimum())
        self.progress_bar.setFormat("当前任务:数据处理   进度:%.2f%%" % p)

    def updateTrainProcessbar(self, x):
        self.progress_bar.setValue(x)
        p = (self.progress_bar.value() - self.progress_bar.minimum()) * 100 / (
                self.progress_bar.maximum() - self.progress_bar.minimum())
        self.progress_bar.setFormat("当前任务:模型训练   进度:%.2f%%" % p)

    def setArgs(self, args: Args):
        args.picture_save_path = self.picture_save_path
        args.model_save_path = self.module_save_path
        args.encoder_type = self.encoder_type
        args.decoder_type = self.decoder_type
        args.dimz = self.dimz
        args.num_epochs = self.epoch
        args.model_name = self.name
        args.batch_size = self.batch_size
        args.learning_rate = self.LR
        args.train_data_path = self.train_data_path
        # args.train_label_path = self.train_label_path

    def updatePlot(self, loss_lst):
        self.test_loss_lst.append(loss_lst[0])
        x_axis = np.arange(1, len(self.test_loss_lst) + 1)
        self.dev_plot.setData(x_axis, self.test_loss_lst)
        self.train_loss_lst.append(loss_lst[1])
        self.train_plot.setData(x_axis, self.train_loss_lst)

    def updateThreshold(self, lst: list):
        print(lst)
        # 1. 找到item
        model_manage = self.parent.setting_dock.model_manage
        for i in range(model_manage.childCount()):
            item: QTreeWidgetItem = model_manage.child(i)
            if item.text(0) == self.band:
                break
        for i in range(item.childCount()):
            sub_item: QTreeWidgetItem = item.child(i)
            if sub_item.text(0) == self.name:
                index = i
                break
        sub_item.setText(1, "%.3f" % (lst[0] + self.n * lst[1]))
        self.parent.models_dict.get(self.band)[self.name]['thresh'] = "%.4f" % (lst[0] + self.n * lst[1])

    @pyqtSlot()
    def on_pushButton_start_train_clicked(self):
        if self.training:
            QMessageBox.warning(self, "提示", "当前正在训练模型, 请训练完毕后再试")
            return
        if self.checkParameters(*self.getTrainParameters()) == False:
            return

        """model_train model"""
        # 检查模型文件夹是否完整
        if not os.path.exists(self.train_data_path):
            QMessageBox.critical(self, "提示", "训练数据不存在!")
            return

        self.printToTextBrowser("初始化中, 开始数据处理......", font_family='微软雅黑', font_size='3')



        nperseg = self.stftsetting['nperseg']
        noverlap = self.stftsetting['noverlap']
        nfft =  self.stftsetting['nfft']
        time_segment = self.stftsetting['time_segment']
        Fs = self.stftsetting['Fs']
        F0 = self.stftsetting['F0']

        train_cnt = self.train_cnt
        bin_data_path = self.train_data_path

        # 计算截取后的时频图维度
        tmp = np.linspace(F0 - Fs / 2, F0 + Fs / 2, int(nfft))
        print(tmp)
        lst = [float(s.strip()) for s in self.band.split(" - ")]
        print(lst)
        print(lst[0] == tmp[0])
        self.begin = np.searchsorted(tmp, lst[0], 'left')
        self.end = np.searchsorted(tmp, lst[1], 'left')

        print("begin:", self.begin, "end:", self.end)

        self.progress_bar.setFormat("当前任务:数据处理   进度:%.2f%%" % 0)
        self.progress_bar.setAlignment(Qt.AlignVCenter | Qt.AlignRight)
        self.progress_bar.setMinimum(0)
        self.progress_bar.setMaximum(train_cnt)
        self.progress_bar.setValue(0)

        bin_file_reader = BinFileReader(self, train_cnt=train_cnt,
                                        bin_file_path=bin_data_path,
                                        noverlap=noverlap,
                                        time_segment=int(time_segment),
                                        nfft=int(nfft),
                                        Fs=Fs,
                                        nperseg=nperseg,
                                        F0=F0,
                                        begin=self.begin,
                                        end=self.end)

        bin_file_reader.requestPassTrainData[object].connect(self.train_model)
        bin_file_reader.requestShowProcess[int].connect(self.updateDealDataProcessbar)
        self.printToTextBrowser("数据处理使用相关参数如下:")

        s = GetHtml(['窗长(nperseg)', '窗函数重叠数(noverlap)', 'FFT长度(nfft)', '时频图时间维度长度', '采样频率(MHz)', '中心频率(MHz)', '完整时频图维度',
                     '截取后维度', '训练时频图数目'],
                    [[nperseg, noverlap, nfft, time_segment, Fs, F0, "(%d, %d)" % (time_segment, nfft),
                      "(%d, %d)" % (time_segment, self.end - self.begin + 1), train_cnt]]
                    )
        self.textBrowser_show_train.append(s)
        bin_file_reader.start()

    def showThreadOutput(self, text):
        self.textBrowser_show_train.append(text)

    def showTrainImage(self, lst):
        """
        :param lst:
        :return:
        """
        input = lst[0]
        rec = lst[1]
        vmin = min(input.min(), rec.min())
        vmax = max(input.max(), rec.max())
        self.widget_input.setImage(input.T, scale=[1, self.draw_ratio], levels=(vmin, vmax))
        self.widget_rec.setImage(rec.T, scale=[1, self.draw_ratio], levels=(vmin, vmax))

    def clearDocks(self):
        try:
            address = self.parent.findChild(DockWidget, name=self.objectName())
            self.parent.needToDeleteDocks.remove(address)
            print("删除了{}窗口".format(self.windowTitle()))
            self.deleteLater()
        except Exception as e:
            print(e)

    @pyqtSlot()
    def on_pushButton_recover_clicked(self):
        self.requestRecoverThread.emit()

    @pyqtSlot()
    def on_pushButton_pause_train_clicked(self):
        self.requestPausedThread.emit()

    @pyqtSlot()
    def on_pushButton_get_model_save_path_clicked(self):
        path = QFileDialog.getExistingDirectory(self, "模型保存路径选取", ".", QFileDialog.ShowDirsOnly)
        if path == "":
            return
        if len(os.listdir(path)) != 0:
            ret = QMessageBox.warning(self, "文件夹非空警告", "选取的文件夹非空, 确认要将模型保存在该路径下吗? 这将清空文件夹下的所有内容。",
                                      QMessageBox.Yes | QMessageBox.No)
            if ret == QMessageBox.No:
                return
            else:
                try:
                    shutil.rmtree(path, ignore_errors=True)
                    os.mkdir(path)
                except Exception as e:
                    QMessageBox.critical(self, "文件清空错误", "{}".format(e))
                    return
        self.lineEdit_model_save_path.setText(os.path.abspath(path))
        self.lineEdit_model_save_path.setCursorPosition(0)

    @pyqtSlot()
    def on_pushButton_get_picture_save_path_clicked(self):
        path = QFileDialog.getExistingDirectory(self, "重构图片保存路径选取", ".", QFileDialog.ShowDirsOnly)
        if path == "":
            return
        if len(os.listdir(path)) != 0:
            ret = QMessageBox.warning(self, "文件夹非空警告", "选取的文件夹非空, 确认要将重构图片保存在该路径下吗? 这将清空文件夹下的所有内容。",
                                      QMessageBox.Yes | QMessageBox.No)
            if ret == QMessageBox.No:
                return
            else:
                try:
                    shutil.rmtree(path, ignore_errors=True)
                    os.mkdir(path)
                except Exception as e:
                    QMessageBox.critical(self, "文件清空错误", "{}".format(e))
                    return
        self.lineEdit_picture_save_path.setText(os.path.abspath(path))
        self.lineEdit_picture_save_path.setCursorPosition(0)

    @pyqtSlot()
    def on_pushButton_get_train_data_path_clicked(self):
        # ('C:/Python Projects/PycharmProjects/未取对数/IQ正常数据/npy文件/train_6400.npy', 'Numpy Files(*.npy)')
        path, _ = QFileDialog.getOpenFileName(self, "训练数据路径选取", ".", "Bin Files(*.bin)")
        if path == "":
            return
        self.lineEdit_train_data_path.setText(os.path.abspath(path))
        self.lineEdit_train_data_path.setCursorPosition(0)

    @pyqtSlot()
    def on_pushButton_end_train_clicked(self):

        print("结束训练")
        try:
            if self.training:
                self.thread.requestInterruption()
            #     address = self.parent.findChild(DockWidget, name=self.objectName())
            #     self.parent.docks.remove(address)
            #     self.parent.needToDeleteDocks.append(address)
            # else:
            #     self.parent.docks.remove(self.parent.findChild(DockWidget, name=self.objectName()))
            #     self.deleteLater()
        except Exception as e:
            print(e)

    def checkParameters(self, epoch, LR, dimz, encoder_type, decoder_type, module_save_path, picture_save_path,
                        batch_size, train_data_path, model_name, train_cnt):
        if epoch <= 0:
            QMessageBox.warning(self, "训练参数设置错误", "训练轮数为0")
            return False
        self.epoch = epoch

        if LR == 0:
            QMessageBox.warning(self, "训练参数设置错误", "学习率为0")
            return False
        self.LR = LR

        self.dimz = dimz
        self.encoder_type = encoder_type
        self.decoder_type = decoder_type

        # if module_save_path == "":
        #     QMessageBox.warning(self, "训练参数设置错误", "模型保存路径未设置")
        #     return False
        # elif len(os.listdir(module_save_path)) != 0:
        #     QMessageBox.warning(self, "训练参数设置错误", "模型保存路径不为空")
        #     return False

        self.module_save_path = module_save_path
        if not os.path.exists(self.module_save_path):
            os.makedirs(self.module_save_path)

        # if picture_save_path == "":
        #     QMessageBox.warning(self, "训练参数设置错误", "图片保存路径未设置")
        #     return False
        # elif len(os.listdir(picture_save_path)) != 0:
        #     QMessageBox.warning(self, "训练参数设置错误", "图片保存路径不为空")
        #     return False
        self.picture_save_path = picture_save_path
        if not os.path.exists(self.picture_save_path):
            os.makedirs(self.picture_save_path)

        if batch_size == 0:
            QMessageBox.warning(self, "训练参数设置错误", "BatchSize为0")
            return False
        self.batch_size = batch_size

        if train_data_path == "":
            QMessageBox.warning(self, "训练参数设置错误", "训练数据路径为空")
            return False
        self.train_data_path = train_data_path
        self.name = model_name
        self.train_cnt = train_cnt
        return True

    def getTrainParameters(self):
        epoch = self.spinBox_epoch.value()
        LR = self.doubleSpinBox_LR.value()
        # dimz = self.spinBox_dimz.value()
        dimz = 50
        encoder_type = "CNN"
        decoder_type = "CNN"
        module_save_path = self.lineEdit_model_save_path.text()
        picture_save_path = self.lineEdit_picture_save_path.text()
        batch_size = self.spinBox_batch_size.value()
        train_data_path = self.lineEdit_train_data_path.text()
        train_cnt = self.spinBox_train_num.value()

        model_name = self.lineEdit_model_name.text()

        if model_name == "":
            model_name = self.name
        self.lineEdit_model_name.setText(model_name)
        self.setWindowTitle("%s : %s" % (self.band, self.name))
        return epoch, LR, dimz, encoder_type, decoder_type, module_save_path, picture_save_path, batch_size, train_data_path, model_name, train_cnt

    def printToTextBrowser(self, text, color='white', font_size='3', font_family='Times New Roman'):
        content = f"""
        <font size={font_size} color={color} style= "font-family:{font_family}">{text}</font>
        """
        self.textBrowser_show_train.append(content)

    def showNotifications(self, lst: list):
        if lst[0] == "info":
            NotificationWindow.info("提示", lst[1])
        elif lst[0] == "warning":
            NotificationWindow.warning("提示", lst[1])
        elif lst[0] == "error":
            NotificationWindow.error("提示", lst[1])
        elif lst[0] == "success":
            NotificationWindow.success("提示", lst[1])

    def train_model(self, train_data):
        print("train_model")
        print("================================")
        print(train_data.shape)
        print("================================")

        # # todo 创建训练进程, 给训练进程绑定槽函数
        self.dev_plot.clear()
        self.train_plot.clear()
        self.test_loss_lst.clear()
        self.train_loss_lst.clear()
        self.updateTrainProcessbar(0)
        args = Args()
        self.setArgs(args)
        self.progress_bar.setMinimum(0)
        self.progress_bar.setMaximum(args.num_epochs)
        self.progress_bar.setFormat("当前任务:模型训练   进度:%.2f%%" % 0)

        tsamples = int(self.parent.setting_dock.spinbox_time_segment.value())
        fsamples = int(self.parent.setting_dock.spinbox_nfft.value())
        nfft = int(self.parent.setting_dock.spinbox_nfft.value())
        noverlap = int(self.parent.setting_dock.spinbox_noverlap.value())
        Fs = float(self.parent.setting_dock.double_spinbox_sample_freq.value() * 1e6)
        F0 = float(self.parent.setting_dock.double_spinbox_center_freq.value() * 1e6)

        self.fsamples = self.end - self.begin + 1
        self.tsamples = tsamples
        self.draw_ratio = int(7 * (self.fsamples / self.tsamples) / 25.6)
        self.thread = TrainProcess(self, args, MultiTrain=self.MultiTrain, train_data_num=self.train_cnt,
                                   train_data=train_data)
        self.thread.updateProcessbar[float].connect(self.updateTrainProcessbar)
        self.thread.updateThreshold[list].connect(self.updateThreshold)
        self.requestPausedThread.connect(self.thread.pause)
        self.requestRecoverThread.connect(self.thread.recover)
        self.thread.updatePlotOutput[list].connect(self.updatePlot)
        self.thread.finished.connect(self.clearDocks)
        self.thread.requestShowPromption[list].connect(self.showNotifications)
        self.thread.requestShowTrainImage[list].connect(self.showTrainImage)
        self.thread.requestAppendToTextBrower[str].connect(self.showThreadOutput)
        self.thread.start()
